import math
from typing import List, Optional, Tuple
import pandas as pd
import numpy as np


class GA_TSP:
    """遗传算法求解TSP问题类"""
    
    def __init__(self, cities: np.ndarray, max_gen: int = 500,
                 pop_size: int = 200, cross_rate: float = 0.8,
                 muta_rate: float = 0.005) -> None:
        self.dist_mat: Optional[np.ndarray] = None
        self.max_gen: int = max_gen  # 最大进化次数
        self.pop_size: int = pop_size  # 种群大小（个体个数）
        self.cross_rate: float = cross_rate  # 交叉概率
        self.muta_rate: float = muta_rate  # 变异概率
        self.cities: np.ndarray = cities  # 城市的坐标数据
        self.cities_num: int = cities.shape[0]  # 城市的个数(个体的染色体长度)
        self.population: List[np.ndarray] = []  # 种群
        self.new_pop: List[np.ndarray] = []  # 新种群
        self.dis: List[float] = []  # 每个个体的距离
        self.fitness: List[float] = []  # 每个个体的适应度
        self.index: int = 0  # 最优个体索引

    # 获得距离矩阵
    def gen_dist_mat(self) -> None:
        D = self.cities
        city_cnt = self.cities_num
        dist_mat = np.zeros((city_cnt, city_cnt))
        for i in range(city_cnt):
            for j in range(city_cnt):
                if i == j:
                    # 相同城市不允许访问
                    dist_mat[i][j] = 1000000
                else:
                    # 单位：km
                    dist_mat[i][j] = 6378.14 * math.acos(
                        math.cos(D[i][1]) * math.cos(D[j][1]) * math.cos(D[i][0] - D[j][0]) +
                        math.sin(D[i][1]) * math.sin(D[j][1]))
        self.dist_mat = dist_mat

    # 计算路径长度
    def dis_cal(self, path: np.ndarray) -> float:
        """
        适应度函数，计算目标函数值.
        :param path: 一个解
        :return: 目标函数值，即总距离
        """
        distance: float = 0.0
        if self.dist_mat is None:
            raise ValueError("距离矩阵未初始化，请先调用gen_dist_mat()方法")
        
        for i in range(len(path) - 1):
            distance += float(self.dist_mat[path[i]][path[i + 1]])
        distance += float(self.dist_mat[path[-1]][path[0]])  # 从最后一个城市回到出发城市
        return distance

    # 初始化种群
    def init_population(self) -> None:
        path: np.ndarray = np.array(range(self.cities_num))
        self.population = []
        for i in range(self.pop_size):
            x: np.ndarray = path.copy()
            np.random.shuffle(x)  # 随机洗牌
            self.population.append(x)

    # 自然选择，种群根据适应度进行更新
    def select(self) -> None:
        # 计算每条路径的长度，放入列表
        self.dis = []
        for i in range(self.pop_size):
            self.dis.append(self.dis_cal(path=self.population[i]))
        # 根据路径长度计算每个个体的适应度
        self.fitness = []
        for i in range(self.pop_size):
            self.fitness.append(1 / (self.dis[i] ** 15))
        # 适应度总和
        fitness_sum: float = 0.0
        for i in range(self.pop_size):
            fitness_sum += self.fitness[i]
        # 根据适应度进行选择，适应度大的被选择概率大
        fitness_probs: np.ndarray = np.array(self.fitness) / fitness_sum
        idx: np.ndarray = np.random.choice(np.arange(self.pop_size), size=self.pop_size, replace=True,
                               p=fitness_probs)
        self.new_pop = []
        for i in idx:
            self.new_pop.append(self.population[i])
        self.population = self.new_pop

    # 交叉互换
    def crossover(self) -> None:
        self.new_pop = []
        for father in self.population:
            child: np.ndarray = father  # 初步让子代获得父本染色体
            if np.random.rand() < self.cross_rate:
                # 随机选择一个染色体作为母本
                mother_index: int = np.random.randint(0, self.pop_size)
                mother: np.ndarray = self.population[mother_index]
                # 确定切割点     
                left: int = np.random.randint(0, self.cities_num - 2)
                right: int = np.random.randint(left + 1, self.cities_num - 1)
                mother_list: List[int] = mother.tolist()
                father_list: List[int] = father.tolist()
                # 切割片段
                gene: List[int] = mother_list[left:right]
                child1_c: List[int] = father_list[right:] + father_list[:right]
                child1: List[int] = child1_c.copy()
                # 去除父本中母本出现过的重复基因
                for o in gene:
                    child1_c.remove(o)
                child1[left:right] = gene # 先填母本基因
                child1[right:] = child1_c[0:len(child1) - right] # 顺序填充父本中不重复的基因
                child1[:left] = child1_c[(len(child1) - right):]
                child = np.array(child1) # 再转为 np 对象

            self.new_pop.append(child)
        self.population = self.new_pop

    # 变异
    def mutation(self) -> None:
        for i in range(self.pop_size):
            if np.random.rand() < self.muta_rate:
                child: np.ndarray = self.population[i]
                u: int = np.random.randint(0, self.cities_num - 2)
                v: int = np.random.randint(u + 1, self.cities_num - 1)
                child_x: np.ndarray = child[u + 1:v]
                # 翻转变异
                child_x = child_x[::-1]
                child = np.concatenate((child[0:u + 1], child_x, child[v:]))
                self.population[i] = child

    # 输出当前种群中最优路径
    def result_path(self) -> None:
        self.index = int(np.argmax(self.fitness))
        a: str = "the shortest path is:"
        for i in range(self.cities_num - 1):
            a += str(self.population[self.index][i])
            a += "-->"
        a += str(self.population[self.index][-1])
        print(a)
        print("the total distance is", self.dis[self.index])


# 主函数进行迭代
def main(cities: np.ndarray) -> None:
    GA: GA_TSP = GA_TSP(cities)
    GA.gen_dist_mat()
    GA.init_population()
    for i in range(GA.max_gen):
        GA.crossover()
        GA.mutation()
        GA.select()
    GA.result_path()


if __name__ == '__main__':
    original_cities: List[List] = [['西宁', 101.74, 36.56],
                       ['兰州', 103.73, 36.03],
                       ['银川', 106.27, 38.47],
                       ['西安', 108.95, 34.27],
                       ['郑州', 113.65, 34.76],
                       ['济南', 117, 36.65],
                       ['石家庄', 114.48, 38.03],
                       ['太原', 112.53, 37.87],
                       ['呼和浩特', 111.65, 40.82],
                       ['北京', 116.407526, 39.90403],
                       ['天津', 117.200983, 39.084158],
                       ['沈阳', 123.38, 41.8],
                       ['长春', 125.35, 43.88],
                       ['哈尔滨', 126.63, 45.75],
                       ['上海', 121.473701, 31.230416],
                       ['杭州', 120.19, 30.26],
                       ['南京', 118.78, 32.04],
                       ['合肥', 117.27, 31.86],
                       ['武汉', 114.31, 30.52],
                       ['长沙', 113, 28.21],
                       ['南昌', 115.89, 28.68],
                       ['福州', 119.3, 26.08],
                       ['台北', 121.3, 25.03],
                       ['香港', 114.173355, 22.320048],
                       ['澳门', 113.54909, 22.198951],
                       ['广州', 113.23, 23.16],
                       ['海口', 110.35, 20.02],
                       ['南宁', 108.33, 22.84],
                       ['贵阳', 106.71, 26.57],
                       ['重庆', 106.551556, 29.563009],
                       ['成都', 104.06, 30.67],
                       ['昆明', 102.73, 25.04],
                       ['拉萨', 91.11, 29.97],
                       ['乌鲁木齐', 87.68, 43.77]]
    df_cities: pd.DataFrame = pd.DataFrame(original_cities, columns=['城市', '经度', '纬度'])
    cities: np.ndarray = df_cities[['经度', '纬度']].values * math.pi / 180

    main(cities)